#!/usr/bin/python

import sys
import math
import getopt
import types

import acgt2wyk
import sample as mysample

def usage():
    print '''Usage : match.py [options] [file]

    Options:
    --------
    -h, --help                   Display this text and exit.
    -d, --debug                  Show debugging info (intermediate matrices).
    -p, --pseudocount            Use pseudocount.
    -c, --scale                  Scale values using Max and Min.
                                 If you don't know what above option is doing don't use it.
    -s, --sample                 Show samples instead of any matrix.
    -w, --wyk                    Show WYK matrix instead of ACGT matrix
    -l, --lowest                 Show lowest scoring sequence above threshold and exit.
    '''
def main(argv):
    try:
        opts, args = getopt.getopt(argv, "hdpcswl", ["help", "debug", "pseudocount", "scale", "sample", "wyk", "lowest"])
    except getopt.GetoptError, detail:
        print "******", detail
        sys.exit("Please specify file having samples.")

    if (len(args) != 1):
        usage()
        sys.exit(3)

    debug = 0
    pseudocount = 0
    scale = 0
    sample = 0
    wyk = 0
    lowest = 0

    for opt, arg in opts:
        if opt in ("-h", "--help"):
            usage()
            sys.exit()
        elif opt in ("-d", "--debug"):
            debug = 1
        elif opt in ("-p", "--pseudocount"):
            pseudocount = 1
        elif opt in ("-c", "--scale"):
            scale = 1
        elif opt in ("-s", "--scale"):
            sample = 1
            #total_population = arg
        elif opt in ("-w", "--wyk"):
            wyk = 1
        elif opt in ("-l", "--lowest"):
            lowest = 1
            #total_population = arg
        else:
            usage()
            sys.exit()    

    # do all the initial processing
    # this program just handles a single file (having alignments) at a time
    lines = init(''.join(args)) 

    if (debug == 1):
        print "\nOccurence Matrix:"
        showMatrix(occurenceMatrix(lines, pseudocount))
    
        print "\nFrequency Matrix:"
        showMatrix(frequencyMatrix(lines, pseudocount))

    weight_mat = matchWeightMatrix(lines, pseudocount, debug, scale)
    wyk_mat = acgt2wyk.acgt2wyk(weight_mat)

    if (debug == 1):
        print "MATCH Weight Matrix:"
        showMatrix(weight_mat)

        print "\nMATCH WYK Weight Matrix:"
        showMatrix(wyk_mat)
    else:
        if (wyk == 1):
            showMatrix(wyk_mat)
        elif (sample == 1):
            total_population = mysample.main(['-t', lowestScore(lines, weight_mat)], weight_mat)
            showSamples(lines, weight_mat, total_population)
        elif (lowest == 1):
            total_population = mysample.main(['-t', lowestScore(lines, weight_mat)], weight_mat)
            showLowest(lines, weight_mat, total_population)
        else:
            showMatrix(weight_mat)

def init(filename):
    # open file having alignments
    try:
        file = open(filename, 'r')
    except IOError:
        print "No such file %s." % filename
        sys.exit(2)
    except:
        print "Unexpected error:", sys.exc_info()[0]
        sys.exit(2)

    lines = file.readlines()

    for i in xrange(len(lines)):
        lines[i] = (lines[i]).strip()

    # return a list cotaining all the lines in given file
    return lines

def log2(x):
    # log2 base logarithm
    
    return (math.log(x) / math.log(2))

def occurenceMatrix(lines, pseudocount):
    # returns a matrix of base occurence at each position
    base_map = {'A': 0, 'C' : 1, 'G' : 2, 'T' : 3, 'a':0, 'c':1, 'g':2, 't':3}

    if (pseudocount == 0):
        matrix = [[0 for x in xrange(len(lines[0]))] for y in xrange(4)]
    else:
        matrix = [[1 for x in xrange(len(lines[0]))] for y in xrange(4)]

    for line in lines:
        for i in xrange(len(line)):
             matrix[base_map[line[i]]][i] += 1

    return matrix

def frequencyMatrix(lines, pseudocount):
    # returns a matrix of frequency of each base at each position
    freq_mat  = occurenceMatrix(lines, pseudocount)

    if (pseudocount == 0):
        for i in xrange(len(freq_mat)):
            for j in xrange(len(freq_mat[0])):
                freq_mat[i][j] /= 1.0 * (len(lines))
    else:
        for i in xrange(len(freq_mat)):
            for j in xrange(len(freq_mat[0])):
                freq_mat[i][j] /= 1.0 * (len(lines)+4)
                

    return freq_mat

def matchWeightMatrix(lines, pseudocount, debug, scale):
    information = [2 for i in xrange(len(lines[0]))]
    
    # return match weight matrix
    weight_mat = frequencyMatrix(lines, pseudocount)

    for j in xrange(len(information)):
        for i in xrange(4):
            if (weight_mat[i][j] > 0):
                information[j] +=  weight_mat[i][j] * log2(weight_mat[i][j])
            else:
                pass # information[j] = 0 for 0 log2(0)

    min = [0 for i in xrange(len(lines[0]))]
    max = [0 for i in xrange(len(lines[0]))]

    for j in xrange(len(min)):
        min[j] = weight_mat[0][j]
        max[j] = weight_mat[0][j]

        for i in xrange(4):
            if weight_mat[i][j] < min[j]:
                min[j] = weight_mat[i][j]
            if weight_mat[i][j] > max[j]:
                max[j] = weight_mat[i][j]
        
    if (scale == 0):
        for i in xrange(len(weight_mat)):
            for j in xrange(len(weight_mat[0])):
                weight_mat[i][j] = information[j] * weight_mat[i][j]
    else:
        for i in xrange(len(weight_mat)):
            for j in xrange(len(weight_mat[0])):
                weight_mat[i][j] = information[j] * ((weight_mat[i][j] - min[j]) / (max[j]-min[j]))

    if debug == 1:
        print "\nInformation:"
        for i in information:
            print '%.2f' % i,
        print '\n'
            
    return weight_mat

def showMatrix(weight_mat):
    for row in weight_mat:
        for col in row:
            if (type(col) == types.IntType):
                print col,
            elif (type(col) == types.FloatType):
                print '%.2f' % col,
            else:
                print "Matrix element type not supported."
                sys.exit(2)
        print
    
def showSamples(lines, weight_mat, total_population):
    total_pop = total_population # all samples having given length
    min_score = score(lines[0], weight_mat)
    
    for line in lines:
        if score(line, weight_mat) < min_score:
            min_score = score(line, weight_mat)

    for line in total_pop:
        if score(line, weight_mat) >= min_score:
            print line

def showLowest(lines, weight_mat, total_population):    
    total_pop = total_population # all samples having given length
    min_score = score( lines[0], weight_mat)

    lowest_score = min_score
    lowest_scoring_seq = lines[0]
    
    for line in lines:
        if score(line, weight_mat) < min_score:
            min_score = score(line, weight_mat)
            lowest_score = min_score
            lowest_scoring_seq = line

    print lowest_scoring_seq

def lowestScore(lines, weight_mat):
    min_score = score(lines[0], weight_mat)

    lowest_score = min_score
    
    for line in lines:
        if score(line, weight_mat) < min_score:
            min_score = score(line, weight_mat)
            lowest_score = min_score

    return lowest_score

def score(line, weight_mat):
    score = 0
    
    base_map = {'a':0, 'c':1, 'g':2, 't':3, 'A':0, 'C':1, 'G':2, 'T':3}

    for i in xrange(len(line)):
        score += weight_mat[base_map[line[i]]][i]

    return score

if __name__ == "__main__":
    main(sys.argv[1:])
